#ifndef RAY_TRACE_PRO_KD_TREE_H
#define RAY_TRACE_PRO_KD_TREE_H

#include "Basic/GxLibBasic.h"
#include "Windows/GxThreads.h"
#include "RayTraceSystem.h"
#include "Debug.h"
//#define DBG_KD_TRAVERSE_STATS

using namespace GxLibBasic;

namespace RayTracePro
{
	const int TraverseStackLength = 48;
	
	class Geometry;

	template<typename T>
	struct MailboxedElement
	{
		GxWinSystem::Threading::SpinLock Lock;
		unsigned int RayId;
		float t;
		DifferentialGeometry Intersection;
		T Element;
	};

	template <typename T>
	class KdNode
	{	
	private:
		unsigned int flag;
	public:
		inline unsigned int GetIsLeaf()
		{
			return flag & (1<<31);
		}
		inline unsigned int GetSplitDimension()
		{
			return flag & 3;
		}
		inline unsigned int GetOffset()
		{
			return flag & 0x7FFFFFFC;
		}
		inline unsigned int GetElementCount()
		{
			return flag & 0x7FFFFFFF;
		}
		inline void SetElementCount(int val)
		{
			flag = val + (flag & 0x80000000);
		}
		inline void SetIsLeaf(int val)
		{
			flag = (val<<31) + (flag & 0x7FFFFFFF);
		}
		inline void SetSplitDimension(int dim)
		{
			flag = dim + (flag & 0xFFFFFFFC);
		}
		inline void SetOffset(int val)
		{
			flag = val + (flag & 0x80000003);
		}
		union
		{
			float SplitPlane;
			MailboxedElement<T> ** Elements;
			MailboxedElement<T> * OneElement;
			//T* Element;
		};
	};

	template <typename T>
	class KdTree
	{
	public:
		Array<MailboxedElement<T>> Elements;
		Array<KdNode<T>, true> Nodes;
		BBox Bounds;
		~KdTree()
		{
			for (int i = 0; i<Nodes.Count(); i++)
				if (Nodes[i].GetIsLeaf() && Nodes[i].GetElementCount() > 1)
					delete [] Nodes[i].Elements;
		}
	};

	struct BBoxEdge
	{
		unsigned int Type : 1; // Start or End
		unsigned int GeometryId : 31;
		float t;
		static const int Start = 0;
		static const int End = 1;
		inline bool operator < (const BBoxEdge & e)
		{
			if (t == e.t)
				return Type<e.Type;
			else
				return t<e.t;
		}
	};

	template<typename T>
	inline void CreateLeafNode(KdTree<T> & tree, KdNode<T> & curNode, int geomCount, int * geoms)
	{
		curNode.SetIsLeaf(1);
		curNode.SetElementCount(geomCount);
		if (geomCount == 1)
		{
			curNode.OneElement = tree.Elements.Buffer() + geoms[0];
		}
		else if (geomCount > 0)
		{
			curNode.Elements = new MailboxedElement<T>*[geomCount];
			for (int g = 0; g<geomCount; g++)
				curNode.Elements[g] = tree.Elements.Buffer() + geoms[g];
		}
		else
		{
			curNode.Elements = 0;
		}
	}

	template<typename T, typename CostEvaluator>
	void ConstructKdTree(KdTree<T> & tree, T* elements, int nodeId, Array<BBox, true> & bboxes,
		BBoxEdge * edges, int * geoms, int geomCount, int * geomLeft, int * geomRight, int depth, int badRefines)
	{
		tree.Elements.SetSize(geomCount);
		for (int i = 0; i < geomCount; i++)
		{
			tree.Elements[i].Element = elements[i];
			tree.Elements[i].RayId = 0;
		}
		if (bboxes.Count() > 0)
		{
			tree.Bounds = bboxes[0];
			for (int i = 1; i < bboxes.Count(); i++)
			{
				tree.Bounds.Union(bboxes[i]);
			}
		}
		
		ConstructKdTreeRec<T, CostEvaluator>(tree, nodeId, bboxes, tree.Bounds, edges, geoms,
			geomCount, geomLeft, geomRight, depth, badRefines);
		printf("%d nodes\n", tree.Nodes.Count());
		//// re-arrange kd tree
		//Array<KdNode<T>, true> nodes;
		//Array<int, true> nodeIds, queue;
		//nodes.SetCapacity(tree.Nodes.Count());
		//int levelStart = 0;
		//nodes.Add(tree.Nodes[0]);
		//nodeIds.Add(0);
		//while (levelStart < nodes.Count())
		//{
		//	queue.Clear();
		//	for (int i = levelStart; i < nodes.Count(); i++)
		//	{
		//		if (nodes[levelStart].GetIsLeaf())
		//		{
		//			continue;
		//		}
		//		queue.Add(nodeIds[levelStart]+1);
		//		queue.Add(nodes[levelStart].GetOffset()/sizeof(KdNode<T>));
		//	}
		//	levelStart = nodes.Count();
		//}
	}

	template<typename T, typename CostEvaluator>
	void ConstructKdTreeRec(KdTree<T> & tree, int nodeId, Array<BBox, true> & bboxes, const BBox & nodeBBox,
		BBoxEdge * edges, int * geoms, int geomCount, int * geomLeft, int * geomRight, int depth, int badRefines)
	{
		tree.Nodes.Add(KdNode<T>());
		#define CUR_NODE tree.Nodes[nodeId]
		if (geomCount <= CostEvaluator::ElementsPerNode || depth == 0)
		{
			CreateLeafNode(tree, CUR_NODE, geomCount, geoms);
			return;
		}
		int axis = 0;
		// Select split axis
		Vec3 nodeSize; Vec3::Subtract(nodeSize, nodeBBox.Max, nodeBBox.Min);
		if (nodeSize.y > nodeSize.x && nodeSize.y > nodeSize.z) axis = 1;
		else if (nodeSize.z > nodeSize.y && nodeSize.z > nodeSize.x) axis = 2;
		int otherAxis[3][2] = {{1, 2}, {0, 2}, {0, 1}};
		int axisRetry = 0;
selectOtherAxis:
		int axis0 = otherAxis[axis][0];
		int axis1 = otherAxis[axis][1];
		// Initialize geometry edges
		int candidates = 0;
		float boundsMin = nodeBBox.Min[axis];
		float boundsMax = nodeBBox.Max[axis];
		for (int i = 0; i<geomCount; i++)
		{
			int geomId = geoms[i];
			edges[i*2].GeometryId = edges[i*2+1].GeometryId = geomId;
			edges[i*2].t = bboxes[geomId].Min[axis]; edges[i*2].Type = BBoxEdge::Start;
			edges[i*2+1].t = bboxes[geomId].Max[axis]; edges[i*2+1].Type = BBoxEdge::End;
			if (edges[i*2].t > boundsMin && edges[i*2].t < boundsMax)
			{
				candidates ++;
			}
			if (edges[i*2+1].t > boundsMin && edges[i*2+1].t < boundsMax)
			{
				candidates ++;
			}
		}
		if (candidates == 0 && axisRetry < 3)
		{
			axis = (axis + 1)% 3;
			axisRetry ++;
			goto selectOtherAxis;
		}
		// Sort edges
		std::sort(edges, edges+geomCount*2);
		// Compute minimum cost split
		float sideArea2 = nodeSize[axis0] * nodeSize[axis1] * 2.0f;
		float invNodeArea = 1.0f/(sideArea2 + 2.0f *nodeSize[axis]* (nodeSize[axis0]+nodeSize[axis1]));
		float noSplitCost = CostEvaluator::IntersectTime * geomCount;
		float minCost = FLT_MAX;
		int bestSplit = -1;
		int nLeft = 0, nRight = geomCount;
		
		for (int i = 0; i<geomCount*2; i++)
		{
			BBoxEdge &thisEdge = edges[i];
			if (thisEdge.Type == BBoxEdge::End)
				nRight --;
			if (thisEdge.t > boundsMin && thisEdge.t < boundsMax)
			{
				float l1 = thisEdge.t - boundsMin;
				float l2 = boundsMax - thisEdge.t;
				float pLeft = (sideArea2 + 2.0f*l1*(nodeSize[axis0]+nodeSize[axis1]))*invNodeArea;
				float pRight = (sideArea2 + 2.0f*l2*(nodeSize[axis0]+nodeSize[axis1]))*invNodeArea;
				float be = (nLeft==0||nRight==0)?CostEvaluator::EmptyNodeBounus:0.0f;
				float cost = CostEvaluator::TraversalTime + (1.0f-be)*CostEvaluator::IntersectTime*(nLeft*pLeft+nRight*pRight);
				if (nodeId == 64 && i == 98)
				{
					printf("98, splitT = %f\n", thisEdge.t);
					printf("cost %f minCost %f, cost < minCost = %d, cost == minCost = %d\n", cost, minCost, (int)(cost<minCost), (int)(cost == minCost));
				}
				if (nodeId == 64 && i == 97)
				{
					printf("97, splitT = %f\n", thisEdge.t);
				}
				if (cost < minCost)
				{
					minCost = cost;
					bestSplit = i;
				}
			}
			if (thisEdge.Type == BBoxEdge::Start)
				nLeft ++;
		}
		if (minCost > noSplitCost && geomCount < 64 * CostEvaluator::ElementsPerNode)
			badRefines ++;
		
		if ((minCost > 4.0f * noSplitCost && geomCount < CostEvaluator::ElementsPerNode) || bestSplit == -1
			|| badRefines == 3 || (nLeft == nRight && nRight == geomCount))
		{
			CreateLeafNode(tree, CUR_NODE, geomCount, geoms);
			return;
		}
		// Classify elements
		float splitT = edges[bestSplit].t;
		if (splitT <= boundsMin + Epsilon)
			splitT += Epsilon;
		if (splitT >= boundsMax - Epsilon)
			splitT -= Epsilon;
		nLeft = nRight = 0;
		for (int i = 0; i<bestSplit; i++)
		{
			if (edges[i].Type == BBoxEdge::Start)
				geomLeft[nLeft++] = edges[i].GeometryId;
		}
		for (int i = bestSplit+1; i<geomCount*2; i++)
		{
			if (edges[i].Type == BBoxEdge::End)
				geomRight[nRight++] = edges[i].GeometryId;
		}
		// Recursive split
		BBox subBBox = nodeBBox;
		CUR_NODE.SetIsLeaf(0);
		CUR_NODE.SetSplitDimension(axis);
		CUR_NODE.SplitPlane = splitT;
		float oriMax = subBBox.Max[axis];
		subBBox.Max[axis] = splitT;
		ConstructKdTreeRec<T, CostEvaluator>(tree, nodeId+1, bboxes, subBBox, edges, geomLeft, nLeft, geomLeft, geomRight+geomCount, depth-1, badRefines);
		CUR_NODE.SetOffset(tree.Nodes.Count() * sizeof(KdNode<T>)); // right child id
		subBBox.Max[axis] = oriMax;
		subBBox.Min[axis] = splitT;
		ConstructKdTreeRec<T, CostEvaluator>(tree, tree.Nodes.Count(), bboxes, subBBox, edges, geomRight, nRight, geomLeft, geomRight+geomCount, depth-1, badRefines);
	}

	template<typename T>
	struct KdTraverseItem
	{
		KdNode<T> * node;
		float tMin, tMax;
	};

	// delegate bool IterateAction(DifferentialGeometry & inter, T * elem, const Ray & ray, const RayDifferential & rayDiffIn, RayDifferential & diffOut, float & t)

	template<typename T, typename Tracer, bool pred, bool useMailbox>
	bool TraverseKdTree(RayTraceSystem * system, const Tracer & tracer, DifferentialGeometry & rs, KdTree<T> & tree, const Ray & ray, const RayDifferential & rayDiff, float & t)
	{
#ifdef DBG_KD_TRAVERSE_STATS
		int nodeTraversed = 0;
		int leafTraversed = 0;
		int emptyLeafTraversed = 0;
		int triangleTraversed = 0;
#endif
		static volatile long rayIdAllocator = 1;
		float tmin, tmax;
		float normalDot = 10000.0f;
		if (!RayBBoxIntersection(tree.Bounds, ray, tmin, tmax))
			return false;
		tmax += Epsilon;
		tmin -= Epsilon;
		bool retVal = false;
		KdTraverseItem<T> stack[TraverseStackLength];
		KdNode<T> * node = tree.Nodes.Buffer();
		Ray traceRay = ray;
		int stackPtr = 0;
		traceRay.tMax = tmax;
		traceRay.tMin = tmin;
		rs.Normal = traceRay.Direction;
		rs.Position.SetZero();
		rs.Geometry = 0;
		unsigned int rayId = 1;
		if (useMailbox)
			rayId = _InterlockedExchangeAdd(&rayIdAllocator, 1);	
		while (node)
		{
			if (traceRay.tMax < tmin)
				break;
			if (tmin > tmax)
				break;
#ifdef DBG_KD_TRAVERSE_STATS
			nodeTraversed ++;
#endif
			if (node->GetIsLeaf())
			{
				DifferentialGeometry inter;
				inter.Init();
				int count = node->GetElementCount();
#ifdef DBG_KD_TRAVERSE_STATS
				leafTraversed ++;
				triangleTraversed += count;
				if (count == 0)
					emptyLeafTraversed ++;
#endif
				MailboxedElement<T>** startElem = node->Elements;
				if (count == 1) startElem = &node->OneElement;
				for (MailboxedElement<T>** elem = startElem; elem < startElem + count; elem++)
				{
					MailboxedElement<T>* curElem = *elem;
					// check mailbox
					bool mailboxHit = false;
					if (useMailbox)
					{
						curElem->Lock.Lock();
						mailboxHit = (curElem->RayId == rayId);
						if (mailboxHit)
						{
							if (curElem->t >= traceRay.tMin - Epsilon && curElem->t <= traceRay.tMax + Epsilon)
							{
								if (pred)
								{
									curElem->Lock.Unlock();
									return true;
								}
								else
									retVal = true;
								rs = inter;
								t = curElem->t;
								traceRay.tMax = curElem->t;
							}
						}
						curElem->Lock.Unlock();
					}
					if (!mailboxHit)
					{
						float it = -1.0f;
						if (tracer.Trace(inter, curElem->Element, traceRay, rayDiff, it))
						{
							if (pred)
								return true;		
							//float ndot = 0.0f;
							if (it <= traceRay.tMax/* - Epsilon || 
								(it < traceRay.tMax + Epsilon && 
								(ndot = Vec3::Dot(inter.Normal, traceRay.Direction)) < normalDot)*/)
							{
								//normalDot = ndot;
								rs = inter;
								traceRay.tMax = it;
								t = it;
								retVal = true;
							}
							if (useMailbox)
							{
								curElem->Lock.Lock();
								curElem->RayId = rayId;
								curElem->t = it;
								curElem->Intersection = inter;
								curElem->Lock.Unlock();
							}
						}
					}
				}
				if (stackPtr)
				{
					KdTraverseItem<T> & item = stack[--stackPtr];
					node = item.node;
					tmax = item.tMax;
					tmin = item.tMin;
				}
				else
					break;
			}
			else // Interior node
			{
				int axis = (int)(node->GetSplitDimension());
				//bool coplanar = abs(ray.Direction[axis]) < Epsilon;
				float tplane;
				//if (!coplanar)
				tplane = (node->SplitPlane - ray.Origin[axis]) * ray.ReciprocalDirection[axis];
				KdNode<T>* node1, *node2;
				float rayOrigin = ray.Origin[axis];
				if (rayOrigin <= node->SplitPlane)
				{
					node1 = node + 1;
					node2 = (KdNode<T>*)((char*)tree.Nodes.Buffer() + node->GetOffset());
				}
				else
				{
					node1 = (KdNode<T>*)((char*)tree.Nodes.Buffer() + node->GetOffset());
					node2 = node + 1;
				}

				if (/*coplanar||*/(tplane > tmax || tplane < 0))
				{
					node = node1;
				}
				else if (tplane < tmin)
				{
					node = node2;
				}
				else
				{
					KdTraverseItem<T> toTraverse;
					toTraverse.node = node2;
					toTraverse.tMin = tplane;
					toTraverse.tMax = tmax;
					stack[stackPtr++] = toTraverse;
					node = node1;
					tmax = tplane;
				}
			}
		}
#ifdef DBG_KD_TRAVERSE_STATS
		if (system->DebugPrint)
		{
			printf("Nodes traversed: %d\n", nodeTraversed);
			printf("leaves traversed: %d\n", leafTraversed);
			printf("empty leaf traversed: %d\n", emptyLeafTraversed);
			printf("triangles traversed: %d\n", triangleTraversed);
		}
#endif
		
		return retVal;
	}

	template <typename T, typename CreateLeafSpheresFunction>
	void ComputeKdTreeBoundingSpheres(KdTree<T> &tree, int depth, Array<Vec4, true> &buffer, CreateLeafSpheresFunction & cleaf)
	{
#define ADD_SPHERE \
			Vec3::Add(position, curBox.Min, curBox.Max); \
			Vec3::Scale(position, position, 0.5f); \
			sphere.x = position.x; sphere.y = position.y; sphere.z = position.z; \
			Vec3::Subtract(range, curBox.Max, curBox.Min); \
			sphere.w = range.Length() * 0.5f; \
			buffer.Add(sphere);
#ifdef DEBUG
		if (depth > 5)
			throw "Depth too large";
#endif
		int queue1[32], queue2[32];
		int *toTraverse = queue1, *nextQueue = queue2;
		int qPtr = 1, nPtr = 0;
		Vec3 position, range;
		Vec4 sphere;
		BBox boxes1[32], boxes2[32];
		BBox *boxes = boxes1, *nextBoxes = boxes2;
		int bboxCount = 0;
		toTraverse[0] = 0;
		boxes[0] = tree.Bounds;
		while (qPtr)
		{
			depth --;
			for (int i = 0; i<qPtr; i++)
			{
				int nodeId = toTraverse[i];
				BBox curBox = boxes[i];
				if (tree.Nodes[nodeId].GetIsLeaf())
				{
#ifdef DEBUG
					if (spheres >= bufferSize)
						throw "Output buffer overflow";
#endif
					cleaf(curBox, nodeId, depth);
				}
				else
				{
					if (depth == 0)
					{
#ifdef DEBUG
						if (spheres >= bufferSize)
							throw "Output buffer overflow";
#endif
						ADD_SPHERE
					}
					else
					{
#ifdef DEBUG
						if (nPtr >= 32)
							throw "Buffer overflow";
#endif
						if (tree.Nodes[nodeId+1].GetElementCount())
						{
							BBox box;
							box = curBox;
							box.Max[tree.Nodes[nodeId].GetSplitDimension()] = tree.Nodes[nodeId].SplitPlane;
							
							nextBoxes[nPtr] = box;
							nextQueue[nPtr] = nodeId+1;
							nPtr++;
						}
						int subId = tree.Nodes[nodeId].GetOffset()>>3;
						if (tree.Nodes[subId].GetElementCount())
						{
							BBox box;
							box = curBox;
							box.Min[tree.Nodes[nodeId].GetSplitDimension()] = tree.Nodes[nodeId].SplitPlane;
							
							nextBoxes[nPtr] = box;
							nextQueue[nPtr] = subId;
							nPtr++;
						}
					}
					
				}
			}
			int * tmpQueue = nextQueue;
			nextQueue = toTraverse;
			toTraverse = tmpQueue;
			BBox * tmpBox = boxes;
			boxes = nextBoxes;
			nextBoxes = tmpBox;
			qPtr = nPtr;
			nPtr = 0;
		}
	}

}

#endif